import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from datetime import datetime
import time
import shutil
from pathlib import Path

class PoseVisualizer:
    """Pose可视化类，用于在原图上显示关键点并连接成多边形"""
    
    def __init__(self):
        # 颜色定义
        self.colors = {
            'left_top': (0, 255, 0),    # 绿色
            'right_top': (0, 0, 255),    # 红色
            'left_bottom': (255, 0, 0),  # 蓝色
            'right_bottom': (0, 255, 255), # 黄色
            'left_0': (255, 0, 255),     # 紫色
            'right_0': (255, 255, 0),    # 青色
            'top_0': (128, 0, 128),      # 深紫色
            'bottom_0': (0, 128, 128)    # 深青色
        }
        
    def visualize_pose(self, pose_dict, image_path, output_dir="results_pose_visualization"):
        """
        在原图上可视化Pose点并自动连接成多边形
        
        Args:
            pose_dict: 姿态点字典，如 {'left_top': (980.0, 180.0), ...}
            image_path: 原始图像路径
            output_dir: 输出目录
            
        Returns:
            保存可视化结果的图像路径
        """
        os.makedirs(output_dir, exist_ok=True)
        image = cv2.imread(image_path)
        if image is None:
            print(f"错误: 无法读取图像 {image_path}")
            return None
        
        filename = os.path.basename(image_path)
        output_path = os.path.join(output_dir, f"pose_{filename}")
        
        # 转换为点集并计算凸包
        points = [((x, y), name) for name, (x, y) in pose_dict.items()]
        if len(points) >= 3:
            pts = np.array([pt for pt, _ in points], dtype=np.float32)
            hull_indices = cv2.convexHull(pts, returnPoints=False).flatten()
            ordered_points = [points[i] for i in hull_indices]
            ordered_points.append(ordered_points[0])  # 闭合多边形
            
            # 绘制多边形和关键点
            image = self._draw_pose(image.copy(), ordered_points)
        
        cv2.imwrite(output_path, image)
        return output_path
    
    def _draw_pose(self, image, ordered_points):
        """绘制关键点和连接线"""
        # 绘制连接线
        for (pt, _), next_pt in zip(ordered_points[:-1], ordered_points[1:]):
            start = tuple(map(int, pt))
            end = tuple(map(int, next_pt[0]))
            cv2.line(image, start, end, (0, 255, 255), 2)
        
        # 绘制关键点
        for (x, y), name in ordered_points[:-1]:
            if name in self.colors:
                color = self.colors[name]
                cv2.circle(image, (int(x), int(y)), 5, color, -1)
                cv2.putText(image, name, (int(x)+10, int(y)), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
        return image

class PositionVisualizer:
    """位置可视化类，用于比较标准位置和测试位置的矩形"""
    
    def __init__(self, seg_model_path, pose_model_path, standard_image_path, rect_size=(0.8, 2.5), scale_factor=100):
        from yoloSeg import YoloSegPose
        self.yolo_seg_pose = YoloSegPose(seg_model_path, pose_model_path)
        self.standard_vector = self.yolo_seg_pose.set_standard_position(standard_image_path)
        self.rect_size = rect_size
        self.scale_factor = scale_factor
        self.output_dir = "results_position_visualization"
        os.makedirs(self.output_dir, exist_ok=True)
        self.pose_visualizer = PoseVisualizer()

    def calculate_position(self, image_path):
        """计算测试位置向量"""
        pose, test_vector = self.yolo_seg_pose.test_one_image(image_path)
        return pose, test_vector
    
    def draw_rectangle(self, vector, ax, label_prefix="Standard", color=(0, 1, 0), thickness=2):
        """绘制矩形位置"""
        if vector is None:
            return
        
        # 计算矩形角点
        center_x, center_y = vector['center'][0] * self.scale_factor, vector['center'][1] * self.scale_factor
        left_vec = (vector['left_vec'][0] * self.scale_factor, vector['left_vec'][1] * self.scale_factor)
        bottom_vec = (vector['bottom_vec'][0] * self.scale_factor, vector['bottom_vec'][1] * self.scale_factor)
        
        # 计算四个角点
        bottom_left = (center_x, center_y)
        bottom_right = (center_x + bottom_vec[0], center_y + bottom_vec[1])
        top_left = (center_x + left_vec[0], center_y + left_vec[1])
        top_right = (top_left[0] + bottom_vec[0], top_left[1] + bottom_vec[1])
        
        # 绘制矩形
        from matplotlib.patches import Polygon
        pts = np.array([bottom_left, bottom_right, top_right, top_left])
        polygon = Polygon(pts, closed=True, edgecolor=color, linewidth=thickness, 
                         fill=False, label=f"{label_prefix}Rect")
        ax.add_patch(polygon)

def process_image_pair(visualizer, color_path, depth_path, filename):
    """处理单个图像对"""
    # 移动文件到处理目录
    processing_dir = "data/processing"
    os.makedirs(os.path.join(processing_dir, "Color"), exist_ok=True)
    os.makedirs(os.path.join(processing_dir, "Depth"), exist_ok=True)
    
    color_processing = os.path.join(processing_dir, "Color", filename)
    depth_processing = os.path.join(processing_dir, "Depth", filename)
    
    shutil.move(color_path, color_processing)
    # shutil.move(depth_path, depth_processing)
    if os.path.exists(depth_path):
        shutil.move(depth_path, depth_processing)
    
    try:
        # 计算位置和姿态
        pose, test_vector = visualizer.calculate_position(color_processing)
        
        # 可视化结果
        fig, ax = plt.subplots(figsize=(8, 8))
        visualizer.draw_rectangle(visualizer.standard_vector, ax, "Standard", (0, 1, 0))
        visualizer.draw_rectangle(test_vector, ax, "Test", (0, 0, 1))
        
        # 设置坐标轴
        max_length = max(visualizer.rect_size) * visualizer.scale_factor * 2
        ax.set_xlim(-50, max_length + 100)
        ax.set_ylim(-50, max_length + 100)
        ax.set_aspect('equal')
        ax.legend()
        
        # 保存位置比较结果
        position_output = os.path.join(visualizer.output_dir, f"position_{filename}")
        fig.savefig(position_output)
        plt.close(fig)
        
        # 保存原始图像上的Pose可视化
        if pose:
            pose_output = visualizer.pose_visualizer.visualize_pose(pose, color_processing)
        
        return position_output, pose_output if pose else None
    finally:
        # 清理处理文件
        if os.path.exists(color_processing):
            os.remove(color_processing)
        if os.path.exists(depth_processing):
            os.remove(depth_processing)

def main():
    # 初始化可视化器
    visualizer = PositionVisualizer(
        seg_model_path="./merge_yolo11n_seg_rknn_model",
        pose_model_path="./merge_yolo11n_pose_rknn_model",
        standard_image_path="data/Color1/440.png"
    )
    
    # 监控目录
    color_dir = "data/Color3"
    depth_dir = "data/Depth3"
    processed_files = set()
    
    plt.ion()  # 交互模式
    
    try:
        while True:
            # 获取最新图像对
            color_files = sorted([f for f in os.listdir(color_dir) if f.endswith('.png')], 
                                key=lambda x: os.path.getmtime(os.path.join(color_dir, x)), 
                                reverse=True)
            
            for filename in color_files:
                if filename in processed_files:
                    continue
                    
                color_path = os.path.join(color_dir, filename)
                depth_path = os.path.join(depth_dir, filename)
                
                if not os.path.exists(depth_path):
                    print(f"缺少深度图: {filename}")
                    continue
                
                print(f"处理: {filename}")
                position_output, pose_output = process_image_pair(visualizer, color_path, depth_path, filename)
                processed_files.add(filename)
                
                # 显示处理进度
                print(f"位置比较结果: {position_output}")
                if pose_output:
                    print(f"Pose可视化结果: {pose_output}")
                
                # time.sleep(0.1)  # 避免过于频繁处理
                break  # 每次只处理一个最新文件
            
            time.sleep(0.5)  # 检查间隔
            
    except KeyboardInterrupt:
        print("用户中断")
    finally:
        plt.ioff()
        plt.close('all')

if __name__ == "__main__":
    main()